# Sarsa Cliff on-policy
import gym
import numpy as np


class SarsaAgent(object):
    def __init__(self, observation_n, act_n, learning_rate=0.01, gamma=0.9, epsilon=0.1):
        self.act_n = act_n  # 动作维度，有几个动作可选
        self.lr = learning_rate  # 学习率
        self.gamma = gamma  # reward的衰减率
        self.epsilon = epsilon  # 按一定概率随机选动作
        self.Q = np.zeros((observation_n, act_n))  # Q表

    # 根据输入观察值，采样输出的动作值，采用epsilon-greedy探索
    def sample(self, observation):
        if np.random.uniform(0, 1) < (1.0 - self.epsilon):  # 根据Q-table的Q值选动作:：当随机数<0.9时，采用策略预测动作
            action = self.predict(observation)  # 采取Q_learning策略预测obs->动作
        else:
            action = np.random.choice(self.act_n)  # 当0.1<随机数<0.9时，采取随机动作
        return action

    # 根据输入观察值，采用Sarsa的方法预测输出的动作值
    def predict(self, observation):
        Q_list = self.Q[observation, :]  # 把Q表截取到对应对应Q_list准备选择动作1
        maxQ = np.max(Q_list)
        action_list = np.where(Q_list == maxQ)[0]  # 找到最大Q值的位置，maxQ可能对应多个action:[n][0]
        action = np.random.choice(action_list)  # 选择对应aciton
        return action

    # Sarsa学习，也就是更新Q-table的方法
    def learn(self, observation, action, reward, next_observation, next_action, done):
        predict_Q = self.Q[observation, action]  # 先将Q表完整复制下来
        if done:  # 当完成游戏后就直接结束，将最终R拷贝下来
            target_Q = reward  # 没有下一个状态了
        else:  # 当游戏未结束时
            target_Q = reward + self.gamma * self.Q[next_observation, next_action]  # Sarsa更新方式
        self.Q[observation, action] += self.lr * (target_Q - predict_Q)  # 一定学习率修正q，用迭代后的Q与迭代前的Q做修正


# 训练与测试
def run_episode(env, agent):
    total_steps = 0  # 记录每个episode走了多少step
    total_reward = 0  # 记录总共获得的奖励
    observation = env.reset()  # 重置环境, 重新开一局（即开始新的一个episode）
    action = agent.sample(observation)  # 根据算法选择一个动作，Sarsa开始确定动作后一直前进
    while True:
        # #注意：用gym时奖励都是通过与环境交互直接获得的，若要修改奖励方式则需要在对应位置手动更正
        next_observation, reward, done, _ = env.step(action)  # 用动作与环境进行一次交互
        next_action = agent.sample(next_observation)  # 根据算法选择一个动作(由下一状态s'得到动作)
        # 用当前s,a,r,s'学习Q-table；训练Sarsa
        agent.learn(observation, action, reward, next_observation, next_action, done)

        action = next_action  # 与Qlearning不同，这里动作通过下一个s'得到，而不是由s得到
        observation = next_observation  # 存储上一个观察值，转移state
        total_reward += reward  # 总训练结束后价值
        total_steps += 1  # 计算step数
        if done:
            break
    return total_reward, total_steps


# 当测试时就不需要迭代Q表，直接预测对应值
def test_episode(env, agent):
    total_reward = 0
    total_steps = 0
    observation = env.reset()
    while True:
        action = agent.predict(observation)  # epsilon-greedy搜索，直接用学到的Q表做探索
        next_observation, reward, done, _ = env.step(action)  # 得到环境信息
        total_reward += reward
        observation = next_observation
        total_steps += 1  # 计算总的测试步数
        if done:
            break
    return total_reward, total_steps


if __name__ == '__main__':
    # 使用gym创建悬崖环境
    env = gym.make("CliffWalking-v0")  # 4个动作：0-up,1-right,2-down,3-left

    # 创建一个agent实例，输入超参数
    agent = SarsaAgent(observation_n=env.observation_space.n, act_n=env.action_space.n,
                       learning_rate=0.1, gamma=0.9, epsilon=0.1)

    # 多次训练，打印每个episode的分数
    for episode in range(250):
        ep_reward, ep_steps = run_episode(env, agent)
        print('Episode %s: steps = %s , reward = %.1f' % (episode, ep_steps, ep_reward))

    # 全部训练结束，查看算法效果
    test_reward, test_steps = test_episode(env, agent)
    print('Sarsa:\ntest steps = %s\ntest reward = %.1f' % (test_steps, test_reward))
